#include <sys/types.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <sys/ioctl.h>
#include <sys/select.h>

#include <unistd.h>
#include <fcntl.h>
#include <signal.h>

#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>

static int ourtty; // Our own tty
static int masterpt;

#define ENUM_SIZE sizeof(compare)/sizeof(char *)

enum argtype {
    ARG_USRNAME = 1,
    ARG_PASSWD,
    ARG_ADDRESS
};

static const char *compare[] = {
    "assword:", // Asking for a password,
    "The authenticity of host ", // Asks to authenticate host
    "(/root/.ssh/id_rsa):",
    "Overwrite (y/n)?",
    "empty for no passphrase):",
    "same passphrase again:",
    "finish"
};

int match(const char *reference, const char *buffer, ssize_t bufsize, int &state) {
    int i;

    for (i = 0; reference[state] != '\0' && i < bufsize; ++i) {
        if (reference[state] == buffer[i]) {
            state++;
        } else {
            state = 0;

            if (reference[state] == buffer[i]) {
                state++;
            }
        }
    }

    return state;
}

int handleoutput(int fd , char *userpasswd) {
    char buffer[1024];
    memset(buffer, 0, 1024);
    int ret = 0;

    int numread = read(fd, buffer, sizeof(buffer));

    printf("switch %s\r\n", buffer);

    for (int i = 0; i < ENUM_SIZE; i++) {
        int state = 0;
        match(compare[i], buffer, numread, state);

        if (state > 0) {
            switch (i) {
                case 0: {
                    write(fd, userpasswd, strlen(userpasswd));
                    write(fd, "\n", 1);
                }
                break;

                case 1:
                    if (compare[i][state] == '\0') {
                        return 2;
                    }

                    break;

                case 3: {
                    write(fd, "y\n", 2);
                }
                break;

                case 2:
                case 4:
                case 5:
                    if (compare[i][state] == '\0') {
                        write(fd, "\n", 2);
                    }

                    break;

                case 6: {
                    ret = 1;
                }
                break;

                default: {
                    write(fd, "\n", 2);
                    break;
                }
            }
        }
    }

    return ret;
}

void window_resize_handler(int signum) {
    struct winsize ttysize; // The size of our tty

    if (ioctl(ourtty, TIOCGWINSZ, &ttysize) == 0) {
        ioctl(masterpt, TIOCSWINSZ, &ttysize);
    }
}

static void show_help() {
    printf("Usage: authorize [username] [userpasswd] [host]\r\n");
}

int main(int argc, char *argv[]) {
    if (4 != argc) {
        show_help();
        return -1;
    }

    masterpt = open("/dev/ptmx", O_RDWR);

    if (-1 == masterpt) {
        printf("open /dev/ptmx failed.\r\n");
        return -1;
    }

    fcntl(masterpt, F_SETFL, O_NONBLOCK);

    if (grantpt(masterpt) != 0) {
        printf("Failed to change pseudo terminal's permission");
    }

    if (unlockpt(masterpt) != 0) {
        printf("Failed to unlock pseudo terminal");
    }

    int ourtty = open("/dev/tty", 0);

    const char *name = ptsname(masterpt);
    int childfd;
    int slavept;

    char *curuser = getlogin();
    char *username = argv[ARG_USRNAME];
    char *userpasswd = argv[ARG_PASSWD];
    char *address = argv[ARG_ADDRESS];
    char pubpath[32] = {0};
    char scpcmd[128] = {0};
    char catcmd[128] = {0};
    char rmcmd[128] = {0};
    sprintf(pubpath, "/%s/.ssh/id_rsa.pub", curuser);
    sprintf(scpcmd, "scp ~/.ssh/id_rsa.pub %s@%s:~/.ssh/id_rsa_tmp", username, address);
    sprintf(catcmd, "ssh -o StrictHostKeyChecking=no %s@%s 'cat ~/.ssh/id_rsa_tmp >> ~/.ssh/authorized_keys'", username, address);
    sprintf(rmcmd, "ssh -o StrictHostKeyChecking=no %s@%s 'rm -rf ~/.ssh/id_rsa_tmp'", username, address);

    if (0 == (childfd = fork())) {
        if (setsid() < 0) {
            perror("child1 setsid error");
            exit(-1);
        }

        slavept = open(name, O_RDWR);

        close(slavept);
        close(masterpt);

        if (-1 == access(pubpath, 0)) {
            printf("file not exist\r\n");
            system("ssh-keygen -t rsa -P '' -f ~/.ssh/id_rsa");
        }

        char *arg[] = {"scp", "~/.ssh/id_rsa.pub", "root@192.168.11.3:/root/.ssh/id_rsa_tmp", NULL};
        execvp(arg[0], arg);
        perror("sshpass ");

        exit(5);

        /*
                system(scpcmd);
                //if(!WIFEXITED(ret))
                {
                    perror("sshpass ");
                    write( masterpt, "finish", 6 );
                    exit(0);
                }

                //int ret = system(catcmd);
                char * arg[] = {"scp","/root/.ssh/id_rsa.pub","root@192.168.11.3:/root/.ssh/id_rsa_tmp",0};
                execvp(arg[0],arg);
                perror("sshpass ");

                //if(ret)
                {
                    write( masterpt, "finish", 6 );
                    exit(0);
                }

                ret = system(rmcmd);
                if(ret)
                {
                    if(errno)
                    {
                        perror("sshpass ");
                    }
                    write( masterpt, "finish", 6 );
                    exit(0);
                }
        */
        write(masterpt, "finish", 6);
        close(masterpt);

        /*
                if(-1 == access("/root/.ssh/id_rsa.pub",0))
                {
                    printf("file not exist\r\n");
                    system("ssh-keygen -t rsa -P '' -f ~/.ssh/id_rsa");
                }

                system("scp ~/.ssh/id_rsa.pub root@192.168.11.3:~/.ssh/id_rsa_tmp");
                system("ssh -o StrictHostKeyChecking=no root@192.168.11.3 'cat ~/.ssh/id_rsa_tmp >> ~/.ssh/authorized_keys'");
                system("ssh -o StrictHostKeyChecking=no root@192.168.11.3 'rm -rf ~/.ssh/id_rsa_tmp'");
                write( masterpt, "finish", 6 );
        */
        return 0;
    } else if (childfd < 0) {
        printf("failed to create child process");
        return -1;
    }

    // We are the parent
    slavept = open("/dev/tty", O_RDWR | O_NOCTTY);

    int status = 0;
    int terminate = 0;
    pid_t wait_id;
    sigset_t sigmask, sigmask_select;

    // Set the signal mask during the select
    sigemptyset(&sigmask_select);

    // And during the regular run
    sigemptyset(&sigmask);
    sigaddset(&sigmask, SIGCHLD);

    sigprocmask(SIG_SETMASK, &sigmask, NULL);

    do {
        if (terminate) {
            wait_id = waitpid(childfd, &status, WNOHANG);
        } else {
            fd_set readfd;
            FD_ZERO(&readfd);
            FD_SET(masterpt, &readfd);

            int selret;
            selret = pselect(masterpt + 1, &readfd, NULL, NULL, NULL, &sigmask_select);

            if (selret > 0) {
                if (FD_ISSET(masterpt, &readfd)) {
                    terminate = handleoutput(masterpt , userpasswd);
                }
            }
        }
    } while (!terminate || wait_id == 0 || (!WIFEXITED(status) && !WIFSIGNALED(status)));

    return 1;
}
